from Fuzzer import Fuzzer
import random
import re

class GrammarFuzzer(Fuzzer):
    
    def __init__(self, equalizeCoverage=False):
        Fuzzer.__init__(self, None, equalizeCoverage)
        self.__reset()
    
    def __reset(self):
        self.__matcher = None
        self.__production = dict()
        self.__non_terminals = set()
        self.__start_terminal = None
        self.__anomalyProperty = 0
        self.__anomalyMethod = []

    # requirement: method has to take the newest regular string
    # and return the modified string
    def addAnomalyMethod(self, method):
        self.__anomalyMethod.append(method)

    def setAnomalyProperty(self, prop):
        assert 0 <= prop <= 1
        self.__anomalyProperty = prop
    
    def setStartTerminal(self, terminal):
        if terminal in self.__non_terminals:
            self.__start_terminal = terminal
        else:
            print "Error: start terminal needs a production"
        
    def addProduction(self, non_terminal, production):
        # in case this non-terminal is new
        if not non_terminal in self.__non_terminals:
            self.__production[non_terminal] = []
            self.__non_terminals.add(non_terminal)
        # add production rule
        lst = self.__production[non_terminal]
        lst.append(production)
        
    def printGrammar(self):
        # print header
        print ''
        print '######################'
        print '# Grammar of Fuzzer  #'
        print '######################'
        # print starting symbol
        print ""
        print "Start symbol:"
        print "    ", repr(self.__start_terminal)
        # print non terminals
        print ""
        print "Non-terminals:"
        for nt in self.__non_terminals:
            print "    ", nt
        # print the different production rules
        print ""
        print "Production rules:"
        for nt in self.__production.keys():
            lst = self.__production[nt]
            derivated = ""
            for i in range(len(lst) - 1):
                derivated += repr(lst[i]) + "| "
            derivated += repr(lst[len(lst) - 1])
            print "    " + repr(nt) + " -> " + derivated

    def getNext(self, length=25, joiner=''):
        assert self.__anomalyMethod
        resultString = self.__start_terminal
        if self.__matcher is None:
            self.__matcher = re.compile('|'.join('(%s)' % nt for nt in self.__non_terminals))
        while True:
            if random.random() < self.__anomalyProperty:  # if isAnomaly 
                resultString = random.choice(self.__anomalyMethod)(resultString)
            m = self.__matcher.search(resultString)
            if not m:
                break
            nt = m.group()
            if len(resultString) < min(length * 10, 200):  # XXX not a very generalizable thing ;)
                replacement = random.choice(self.__production[nt])
            else:
                replacement = min(self.__production[nt], key=len)
            resultString = resultString.replace(nt, replacement, 1)
        # print '%d ==> %s' % (length, resultString)
        return resultString
